# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

from transformers.models.llama import LlamaModel
from transformers.models.qwen2 import Qwen2Model
from transformers.models.mistral import MistralModel

from . import qwen2, llama, mistral, internlm2, internlm3

SUPPORTED_FA_MODELS = {
    LlamaModel.__name__: llama,
    Qwen2Model.__name__: qwen2,
    MistralModel.__name__: mistral,
    "InternLM2Model": internlm2,
    "InternLM3Model": internlm3,
}

__all__ = [
    "internlm2",
    "internlm3",
    "llama",
    "qwen2",
    "mistral",
    "SUPPORTED_FA_MODELS",
]
